package datastructure.tree.huffmancode;

import java.util.*;

/**
 * 赫夫曼编码
 *
 * @author jack.wu
 * @version 1.0
 * @date 2020-03-30
 */
public class HuffmanCode {

    public static void main(String[] args) {

        // 压缩
        String content = "i like like like java do you like a java";
        byte[] contentBytes = content.getBytes();
        byte[] zipByte = huffmanZip(contentBytes);
        System.out.println(Arrays.toString(zipByte));

        // 解压
        byte[] decode = decode(huffmanCodes, zipByte);
//        System.out.println(Arrays.toString(decode));
    }

    /**
     * 赫夫曼压缩
     *
     * @param b 原始字符串的字节数组
     * @return 经过赫夫曼编码的压缩字节数组
     */
    private static byte[] huffmanZip(byte[] b) {
        // 得到所有结点
        List<Node> nodes = getNodes(b);
        // 生产赫夫曼树
        Node node = creatHuffmanTree(nodes);
        // 得到赫夫曼编码
        Map<Byte, String> codes = getCodes(node);
        // 压缩
        return zip(b, codes);
    }

    /**
     * 创建赫夫曼树
     *
     * @param nodes
     * @return
     */
    private static Node creatHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            // 排序 从小到大
            Collections.sort(nodes);
            // 取出前两个
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);

            // 创建父节点 没有data
            Node parentNode = new Node(null, leftNode.weight + rightNode.weight);
            parentNode.leftNode = leftNode;
            parentNode.rightNode = rightNode;

            // 移除前两个
            nodes.remove(leftNode);
            nodes.remove(rightNode);

            // 添加到nodes
            nodes.add(parentNode);
        }
        // 返回根节点
        return nodes.get(0);
    }


    /**
     * 获取结点集合
     *
     * @param bytes
     * @return
     */
    private static List<Node> getNodes(byte[] bytes) {
        List<Node> nodes = new ArrayList<>();
        Map<Byte, Integer> counts = new HashMap<>(8);
        for (byte b : bytes) {
            Integer val = counts.get(b);
            if (val != null) {
                counts.put(b, ++val);
            } else {
                counts.put(b, 1);
            }
        }

        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }

    private static Map<Byte, String> huffmanCodes = new HashMap<>(8);
    private static StringBuilder sb = new StringBuilder();

    /**
     * 重载 方便调用
     *
     * @param root
     * @return
     */
    private static Map<Byte, String> getCodes(Node root) {
        if (root == null) {
            return null;
        }
        getCodes(root.leftNode, "0", sb);
        getCodes(root.rightNode, "1", sb);
        return huffmanCodes;
    }

    /**
     * 得到将传入的node结点的所有叶子结点的赫夫曼编码，并放入huffmanCodes
     *
     * @param node 传入的结点
     * @param code 路径：0：左边 1：右边
     * @param sb   拼接路径
     */
    private static void getCodes(Node node, String code, StringBuilder sb) {
        StringBuilder builder = new StringBuilder(sb);
        builder.append(code);
        if (node != null) {
            if (node.data == null) {
                // 说明是非叶子结点 左右递归
                getCodes(node.leftNode, "0", builder);
                getCodes(node.rightNode, "1", builder);
            } else {
                // 说明是叶子结点
                huffmanCodes.put(node.data, builder.toString());
            }
        }
    }

    /**
     * @param bytes        原始字符串对应的byte[]
     * @param huffmanCodes 赫夫曼map
     */
    private static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(huffmanCodes.get(b));
        }
        System.out.println(sb.toString());

        // 字节长度  len = (sb.length() + 7) /8;
        int len;
        if (sb.length() % 8 == 0) {
            len = sb.length() / 8;
        } else {
            len = sb.length() / 8 + 1;
        }
        // 压缩后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        // 记录第几个byte
        int index = 0;
        // 步长8 一个字节8位
        for (int i = 0; i < sb.length(); i += 8) {
            String strByte;
            if (i + 8 > sb.length()) {
                strByte = sb.substring(i);
            } else {
                strByte = sb.substring(i, i + 8);
            }
            huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte);
            index++;
        }

        return huffmanCodeBytes;
    }


    // 解码部分

    /**
     * 传入一个byte 转成一个二进制的字符串
     *
     * @param flg 标志是否需要补高位
     * @param b   压缩过的字节
     * @return 压缩过的字节数组对应的二进制字符串
     */
    private static String byteToBitString(boolean flg, byte b) {
        int temp = b;
        if (flg) {
            // 需要补高位  或运算 如 1 00000000 | 00000001 = 1 00000001
            temp |= 256;
        }
        // 转换二进制字符串
        String str = Integer.toBinaryString(temp);
        if (flg) {
            // 由于转换之后的二进制字符串长度是根据计算机处理器长度来的 而我们只需要后8位
            return str.substring(str.length() - 8);
        } else {
            // 没补高位直接返回
            return str;
        }
    }

    /**
     * 解码
     *
     * @param huffmanCodes 赫夫曼编码表
     * @param huffmanBytes 赫夫曼编码得到的字节数组
     * @return 原来字符串对应的字节数组
     */
    private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            byte b = huffmanBytes[i];
            // 是否是最后一个
            boolean flag = i == huffmanBytes.length - 1;
            // 最后一位不用补高位 因为我们在压缩的时候没有补
            sb.append(byteToBitString(!flag, b));
        }
        // 因为要反向查询 需要把huffmanCodes key和value 交换
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }
        List<Byte> list = new ArrayList<>();
        // 挨个挨个去遍历赫夫曼编码的二进制字符串
        for (int i = 0; i < sb.length(); ) {
            int count = 1;
            boolean flag = true;
            Byte b = null;
            while (flag) {
                String key = sb.substring(i, i + count);
                b = map.get(key);
                if (b == null) {
                    // 没找到，继续向下找
                    count++;
                } else {
                    // 找到退出
                    flag = false;
                }
            }
            list.add(b);
            // i直接移动到count
            i += count;
        }
        byte[] decode = new byte[list.size()];
        for (int i = 0; i < list.size(); i++) {
            decode[i] = list.get(i);
        }
        return decode;
    }

}


class Node implements Comparable<Node> {

    /**
     * 存放数据
     */
    Byte data;
    /**
     * 权值（出现次数）
     */
    int weight;
    Node leftNode;
    Node rightNode;

    public Node(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public int compareTo(Node o) {
        return this.weight - o.weight;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    /**
     * 前序打印
     */
    public void preNode() {
        System.out.println(this);
        if (this.leftNode != null) {
            this.leftNode.preNode();
        }
        if (this.rightNode != null) {
            this.rightNode.preNode();
        }
    }
}